import json
import os

from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field


# 定义结构化输出的模型
class QuestionAnswer(BaseModel):
    thinking: str = Field(description="详细的思考过程")  # 修复字段名称，从"thingking"改为"thinking"
    answer: str = Field(description="思考后回复的答案")

# 定义提示模板
prompt_template = PromptTemplate(
    input_variables=["question"],
    template="请回答以下问题：\n\n问题: {question}\n\n答案:"
)

def get_structured_output(qua,prt:PromptTemplate,question) -> QuestionAnswer:
    #  安全地获取 API 密钥
    os.environ["OPENAI_API_KEY"] = 'local-qwen2.5-72b-little-brother'
    #  设置 base_url 和 model
    base_url = "http://10.8.50.33:8814/v1"  # 替换为你的自定义 base URL
    model = "qwq-32b-preview"

    #  创建 ChatOpenAI 实例
    llm = ChatOpenAI(base_url=base_url, model=model, api_key=os.environ["OPENAI_API_KEY"],max_tokens=18000)
    # 创建带有结构化输出的链
    qa_chain = prt | llm.with_structured_output(qua)

    # 调用链并获取结构化输出
    return  qa_chain.invoke({"question": question})

response = get_structured_output(QuestionAnswer,prompt_template,"什么是龙鳞马？")
# 打印结构化输出为 JSON 格式
print(json.dumps(response.model_dump(), ensure_ascii=False, indent=4))